import pdb
from mpi4py import MPI
from numpy import dot,array
from v3dsolver.vector import Vec
from v3dsolver.fluid_elements.fluid_elems import BoundaryBlob,VortexBlob
from v3dsolver.fluid_elements.fluid_elems import StretchStensilPoint
from v3dsolver.tree_manager.tree_manager import TreeManager
from v3dsolver.tree_manager.vel_calculator import VelCalculator
from v3dsolver.tree_manager.distribution_manager import get_domain

class StretchingManager:
    """
    Class representing the stretching manager object. This will be responsible
    for performing stretching of vorticity on different elements.
    """
    def __init__(self,elems,dx,dt,freestream,npmax,npseudo,blob_rad,
                 vel_method='s'):
        """
        Initialization of StretchingManager object. Parameters as defined below:
        
        elems: total particles present in the simulation, including boundary
               elements
        dx: incremental distance from an element along each axis at which the 
            corresponding stretch stensil point would be placed.
        dt: timestep for vorticity upgradation
        freestream: FreeStream object
        npmax: splitting criterion for tree
        npseudo: number of pseudo particles to generate per cell
        blob_rad: radius of blobs
        vel_method: velocity computation method to be used for computation:
                    'd': direct computation
                    's': serial tree based computation
                    'p': parallel tree based computation
        """
        # boundary elems and parent elems hosting the stretch stensil points
        self.boundaryelems = []
        self.parentelems = []
        for elem in elems:
            if isinstance(elem,BoundaryBlob):
                self.boundaryelems += [elem]
            else:
                self.parentelems += [elem]
        self.dx = self.dy = self.dz = dx
        self.dt = dt
        self.npmax = npmax
        self.npseudo = npseudo
        self.blob_rad = blob_rad
        self.vel_method = vel_method
        self.freestream = freestream

    def set_stretchpts(self):
        """
        Sets stretch stensil points corresponding to elems present in the
        StretchingManager object.
        """
        for elem in self.parentelems:
            if isinstance(elem,VortexBlob):
                stretchpts = []
                # stretchpts in the order: [+x,+y,+z]
                dirs = [Vec(1,0,0),Vec(0,1,0),Vec(0,0,1)]
                for direction in dirs:
                    stretchpts += [StretchStensilPoint(elem,direction,self.dx)]
                elem.stretchpts = stretchpts
    
    def get_stretchpts(self):
        """
        Returns all the stretch stensil points, for the current 
        StretchingManager object.
        """
        stretchpts = []
        for parentelem in self.parentelems:
                stretchpts += parentelem.stretchpts
        return stretchpts
    
    def compute_vels(self):
        """
        Compute velocities on all the stretch stensil points using the method
        described by vel_method attribute of StretchingManager object.
        """
        # direct computation
        if self.vel_method=='d':
            VelCalculator().direct_vel_computation(self.elems)
        # serial tree computation
        elif self.vel_method=='s':
            domain = get_domain(self.elems)
            tree = TreeManager(self.elems,domain,self.npmax,self.npseudo,
                               self.blob_rad,comm=MPI.COMM_WORLD)
            VelCalculator(tree).serial_vel_computation()
        # parallel tree computation
        elif self.vel_method=='p':
            domain = get_domain(self.elems,findglobal=True)
            tree = TreeManager(self.elems,domain,self.npmax,self.npseudo,
                               self.blob_rad,comm=MPI.COMM_WORLD)
            VelCalculator(tree).parallel_vel_computation()
        # effect of freestream
        for elem in self.elems:
            elem.vel += self.freestream.get_vel(elem)

    def update_vors(self):
        """
        Updates the vorticity of each parent elem of StretchingManager 
        according to the stretching equation: dw/dt = w.grad(v)
        
        Note: Euler integration is performed to upgrade vorticity.
        """
        self.set_stretchpts()
        self.stretchpts = self.get_stretchpts()
        self.elems = self.boundaryelems + self.parentelems + self.stretchpts
        self.compute_vels()
        for i,elem in enumerate(self.parentelems):
            # don't stretch boundary blobs
            if not isinstance(elem,BoundaryBlob):
                dv_dx = ((elem.stretchpts[0].vel-elem.vel)/self.dx)
                dv_dy = ((elem.stretchpts[1].vel-elem.vel)/self.dy)
                dv_dz = ((elem.stretchpts[2].vel-elem.vel)/self.dz)
                tmp = Vec()
                for j in [0,1,2]:
                    tmp[j] = self.dt*elem.vor.dot(Vec(dv_dx[j],dv_dy[j],dv_dz[j]))
                elem.vor = elem.vor + tmp
